#!/usr/bin/env python
# -*- coding: utf-8 -*-

from fastapi import FastAPI, WebSocket, Request, Cookie, Query, status, Depends
from fastapi.responses import HTMLResponse
from fastapi.templating import Jinja2Templates
from typing import Optional




app = FastAPI(debug=True)
templates = Jinja2Templates(directory="./templates")




@app.get("/")
async def home(request: Request):
    return templates.TemplateResponse(
        "webchat.html",
        {
            "request": request
        }
    )




async def get_cookie_or_token(
        websocket: WebSocket,
        session: Optional[str] = Cookie(None),
        token: Optional[str] = Query(None),
):
    if session is None and token is None:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
    return session or token


@app.websocket("/item/ws")
async def websocket_endpoint(
        websocket: WebSocket,
        q: Optional[int] = None,
        cookie_or_token: str = Depends(get_cookie_or_token),
):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        if data == 'bye':
            await websocket.send_text("请注意：窗口即将关闭！")
            await websocket.send_text("请注意：窗口已关闭！")
            await websocket.close()
            break
        else:
            await websocket.send_text(
                f"Session cookie or query token value is: {cookie_or_token}"
            )
            if q is not None:
                await websocket.send_text(f"parameter q is: {q}")
            await websocket.send_text(f"Message text was: {data}, for : {q}")




if __name__ == '__main__':
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)